package cn.kinoko.common.utils;

import lombok.extern.slf4j.Slf4j;

import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.LockSupport;

/**
 * @author kinoko
 */
@Slf4j
public class SlidingWindowRateLimiter {
    // 最大请求数
    private final int maxRequest;
    // 时间窗口大小（毫秒）
    private final long windowSizeInMs;
    // 等待时间
    private final long waitTime;
    // 使用并发跳表映射存储每个时间戳对应的请求数，保证线程安全且有序
    private final ConcurrentSkipListMap<Long, AtomicLong> requests = new ConcurrentSkipListMap<>();

    public SlidingWindowRateLimiter(int maxRequest, long windowSizeInMs, long waitTime) {
        this.maxRequest = maxRequest;
        this.windowSizeInMs = windowSizeInMs;
        this.waitTime = waitTime;
    }

    /**
     * 判断当前请求是否被允许。
     * 使用CAS操作减少锁竞争，通过自旋等待机制优化性能，直到窗口内请求数量低于阈值。
     *
     * @param canRetry 是否允许重试
     * @return 如果请求被允许则返回true，否则返回false
     */
    public boolean allowRequest(boolean canRetry) {
        long currentTime = System.currentTimeMillis();

        // 获取头节点到指定key节点之间的所有节点映射
        // 即移除窗口之外的旧请求，ConcurrentSkipListMap的headMap和clear操作天然线程安全
        requests.headMap(currentTime - windowSizeInMs, false).clear();

        // 获取当前时间戳对应的请求数，如果不存在则创建并初始化为0
        Long currentKey = requests.floorKey(currentTime);
        AtomicLong requestCount;
        if (currentKey == null) {
            // 如果当前时间戳对应的请求数不存在，则创建并初始化为1，并且保证只有一个线程参与构建
            AtomicLong initCount = new AtomicLong(1);
            requestCount = requests.putIfAbsent(currentTime, initCount);
            // 构建的那个线程获取到的值将为null
            if (requestCount == null) {
                // 此时创建线程直接放行，防止需要重新获取值，导致被后续等待构建的线程将值变动，使构建线程吃亏
                return true;
            }
        } else {
            requestCount = requests.get(currentKey);
        }
        // 判断请求是否被允许
        if (requestCount.incrementAndGet() <= maxRequest) {
            return true;
        } else {
            // 如果不允许重试，则直接返回false
            if (!canRetry) return false;
            // 睡眠一段时间，重试
            LockSupport.parkNanos(Duration.ofMillis(waitTime).toNanos());
            return allowRequest(false);
        }
    }

    public static void main(String[] args) throws InterruptedException {
        SlidingWindowRateLimiter limiter = new SlidingWindowRateLimiter(1, 1000, 1000); // 示例配置：每秒最多1个请求
        CountDownLatch latch = new CountDownLatch(6);
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            CompletableFuture.runAsync(() -> {
                try {
                    log.info("Request {} allowed: {}", finalI + 1, limiter.allowRequest(true));
                } finally {
                    latch.countDown();
                }
            });
        }
        latch.await();
    }
}
